from torchline.models import META_ARCH_REGISTRY

import torch
import torch.nn as nn

__all__ = [
    'SimpleLSTM'
]

@META_ARCH_REGISTRY.register()
def SimpleLSTM(cfg):
    in_dim = cfg.model.lstm.in_dim
    n_layer = cfg.model.lstm.n_layer
    hidden_dim = cfg.model.lstm.hidden_dim
    classes = cfg.model.lstm.classes
    return _SimpleLSTM(in_dim, hidden_dim, n_layer, classes)

class _SimpleLSTM(nn.Module):
    def __init__(self, in_dim, hidden_dim, n_layer, classes):
        super(_SimpleLSTM, self).__init__()
        self.n_layer = n_layer
        self.hidden_dim = hidden_dim
        self.lstm = nn.LSTM(in_dim, hidden_dim, n_layer, batch_first=True)
        self.classifier = nn.Sequential(
            # nn.Linear(hidden_dim, hidden_dim, bias=False),
            # nn.LeakyReLU(),
            nn.Dropout(0.5),
            nn.Linear(hidden_dim, classes, bias=False)
        )

    def forward(self, x):
        # h0 = Variable(torch.zeros(self.n_layer, x.size(1),
        #   self.hidden_dim)).cuda()
        # c0 = Variable(torch.zeros(self.n_layer, x.size(1),
        #   self.hidden_dim)).cuda()
        out, _ = self.lstm(x)
        out = out[:, -1, :]
        out = self.classifier(out)
        return out